package com.fan.chat.listener;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.fan.chat.model.JSONResult;
import com.plexpt.chatgpt.listener.SseStreamListener;
import com.plexpt.chatgpt.util.SseHelper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

@Slf4j
public class CusSseStreamListener extends SseStreamListener {

    private SseEmitter sseEmitter;

    public CusSseStreamListener(SseEmitter sseEmitter) {
        super(sseEmitter);
        this.sseEmitter = sseEmitter;
        super.setOnComplate(msg -> {
            log.info("request response msg: {}", msg);
            SseHelper.send(this.sseEmitter, doneResult());
            SseHelper.complete(sseEmitter);
        });
    }

    @Override
    public void onMsg(String message) {
        SseHelper.send(this.sseEmitter, JSONResult.success(message));
    }

    @Override
    public void onError(Throwable throwable, String response) {
        log.info("response error: {}", response);
        try {
            JSONObject json = JSON.parseObject(response);
            response = "```json\n" + JSON.toJSONString(json, SerializerFeature.PrettyFormat) + "\n```";
        } catch (Exception ignore) {}
        SseHelper.send(this.sseEmitter, JSONResult.success(response));
        SseHelper.send(this.sseEmitter, doneResult());
        sseEmitter.completeWithError(throwable);
    }

    private static JSONResult<String> doneResult() {
        return JSONResult.success("[DONE]");
    }

}
